# 仿真器入口 用于训练强化学习 跟踪算法用的简化版本仿真器
# 由于没有GUI界面 
# 航控需要通过 位置偏差 和 航向偏差 等指标查看控制算法是否有效 
from dynamics_simulation import motionSim
import numpy as np
import random
import socket
import json
import traceback
from algo_struct.algo import rad_limit

# 解析控制参数
def decode_data(data):
    # load json
    recv_obser = json.loads(data)
    # print(recv_obser)
    delta = recv_obser['rudl']
    rsp = recv_obser['rspl']

    return delta, rsp

# 打包船的状态返回
def encode_data(trans):
    # 按照标准接口编写
    data = json.dumps(trans, sort_keys=True, indent=4, separators=(',', ':'))
    return data.encode('utf-8')

# 计算奖励
def get_reward(observation: dict):
    dx = observation['gx'] - observation['cx']
    dy = observation['gy'] - observation['cy']
    if np.sqrt(dx**2 + dy**2) < 0.5:
        sparse = 1
    else:
        sparse = 0
    # 到达目标点就是奖励
    score = sparse -np.sqrt( dx**2+ dy**2)

    print(f'score is {score}')
    return score

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--port', type=int, default=10064)
args = parser.parse_args()

# test motion sim
if __name__ == '__main__':
    boat_data = dict()
    # 与航行控制软件连接TCP
    LOCAL_HOST = "127.0.0.1"
    LOCAL_PORT = args.port
    
    socket_server = socket.socket(family=socket.AF_INET, type=socket.SOCK_STREAM)
    socket_server.bind((LOCAL_HOST, LOCAL_PORT))
    socket_server.listen()
    conn, address = socket_server.accept()
    print(f"收到了客户端的连接，客户端信息是 {address}")

    all_sim_time = 300
    sim_time = 0
    delta_h = 0.1
    
    delta = 0
    try:
        exp_x = random.randint(50, 100)
        exp_y = random.randint(50, 100)
        exp_angle = random.randint(0, 90)
        x = np.array([0.0, 0.0, 0.0])# 初始无人艇位置
        u = np.array([0.0, 0.0, 0.0])# 初始无人艇速度
        last_u = 0.0
        # 规划进出码头路线 自动泊车路线

        while True:
            sim_time += delta_h
            print(f'期望x和y和angle {exp_x}, {exp_y}, {exp_angle}')
            # 接收转速和舵角
            info, address = conn.recvfrom(1024)
            angular, speed = decode_data(info)
            print(f'速度和方向变化量 {speed}, {angular}')

            # 计算更新 位置 速度
            u[0] += speed # 速度m/s
            
            if u[0] > 10:
                u[0] = 10
            elif u[0]   < -5:
                u[0] = -5
            
            u[2] += angular * 5 * np.pi/180.0# rad
            
            if u[2] > 20 * np.pi/180.0:
                u[2] = 20 * np.pi/180.0
            elif u[2]   < -20 * np.pi/180.0:
                u[2] = -20 * np.pi/180.0
                
            x[2] += u[2] * delta_h # NED-z rad -pi~pi
            x[2] = rad_limit(x[2])
            
            x[0] += u[0] * np.cos(x[2]) * delta_h# NED-x
            x[1] += u[0] * np.sin(x[2]) * delta_h# NED-y
            
            # 为强化学习用
            observation=dict()
            observation['gx'] = exp_x  # x位置偏差
            observation['gy'] = exp_y  # y位置偏差
            observation['cx'] = x[0] # x位置偏差
            observation['cy'] = x[1] # y位置偏差
            observation['psi'] = rad_limit( rad_limit(x[2]) - exp_angle*np.pi/180.0 ) * 180.0/np.pi# rad偏差
            observation['dx'] = exp_x - x[0] # x位置偏差
            observation['dy'] = exp_y - x[1] # y位置偏差
            observation['v_x'] = u[0] * np.cos(x[2]) # x速度
            observation['v_y'] = u[0] * np.sin(x[2]) # y速度
            observation['speed'] = u[0] # 船体坐标系下的轴向速度
            observation['angular'] = u[2] # 角速度
            observation['hgd'] = x[2] # 角度
            observation['target_angle'] = np.arctan2(exp_y - x[1], exp_x - x[0]) # 目标角度 
            boat_data['reward'] = get_reward(observation)
            print(f'observation is {observation}')
            
            if sim_time > all_sim_time or (x[0] > np.abs(200) or x[1] > np.abs(200)):
                exp_x = random.randint(50, 100)
                exp_y = random.randint(50, 100)
                exp_angle = random.randint(0, 90)
                x = np.array([0.0, 0.0, 0.0])# 初始无人艇位置
                u = np.array([0.0, 0.0, 0.0])# 初始无人艇速度

                sim_time = 0
                done = 1
            else:
                done = 0

            boat_data['observation'] = observation
            boat_data['terminated'] = done
            boat_data['truncated'] = done
            data = encode_data(trans=boat_data)
            # 发送数据
            conn.send(data)
            
    except Exception:
        # 关闭连接
        conn.close()
        socket_server.close()
        traceback.print_exc()
